package org.feng.controller;

import jakarta.annotation.Resource;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.document.Document;
import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.advisor.RetrievalAugmentationAdvisor;
import org.springframework.ai.rag.generation.augmentation.ContextualQueryAugmenter;
import org.springframework.ai.rag.postretrieval.document.DocumentPostProcessor;
import org.springframework.ai.rag.preretrieval.query.expansion.MultiQueryExpander;
import org.springframework.ai.rag.retrieval.join.ConcatenationDocumentJoiner;
import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Controller;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;

import java.util.Collections;
import java.util.List;
import java.util.stream.Collectors;

/**
 * 聊天控制器
 *
 * @author pine
 * @version v1.0
 * @since 2025-08-12 20:28
 */
@Controller
@Slf4j
@RequestMapping("/api")
public class ChatController {
    @Resource
    private ChatClient chatClient;
    @Resource
    private ChatMemory chatMemory;
    @Resource
    private VectorStore vectorStore;

    private final ChatClient.Builder chatBuilder;

    public ChatController(ChatClient.Builder chatBuilder) {
        this.chatBuilder = chatBuilder;
    }


    /**
     * 聊天
     */
    @GetMapping(value = "/chat", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    @SneakyThrows
    @ResponseBody
    public Flux<String> chat(@RequestParam(name="conversationId", required = false, defaultValue = "feng123") String conversationId, @RequestParam(name="message", defaultValue = "你好，你有什么功能") String message) {
        log.info("chat(), conversationId:{},message:{}", conversationId, message);

        // 查询扩展
        MultiQueryExpander queryExpander = MultiQueryExpander.builder()
                // 扩展数量：3
                .numberOfQueries(3)
                .chatClientBuilder(chatBuilder)
                // 包括源查询条件
                .includeOriginal(true)
                .build();

        // 增加上下文的信息
        ContextualQueryAugmenter contextualQueryAugmenter = ContextualQueryAugmenter.builder()
                .allowEmptyContext(true)
                .build();


        // RAG
        RetrievalAugmentationAdvisor retrievalAugmentationAdvisor = RetrievalAugmentationAdvisor.builder()
                // 使用向量数据库作为文档源
                .documentRetriever(VectorStoreDocumentRetriever.builder().vectorStore(vectorStore).build())
                // 扩充查询条件
                .queryExpander(queryExpander)
                // 拼接查询到的文档
                .documentJoiner(new ConcatenationDocumentJoiner())
                // 取多个doc中的第一个
                .documentPostProcessors(new SelectedFirstDocumentPostProcessor())
                // 对生成的查询增强上下文
                .queryAugmenter(contextualQueryAugmenter)
                .build();


        return chatClient.prompt(message)
                // 日志
                .advisors(new SimpleLoggerAdvisor())
                // 聊天记忆(暂时使用内存，后续支持jdbc)
                .advisors(MessageChatMemoryAdvisor.builder(chatMemory).conversationId(conversationId).build())
                // RAG
                .advisors(retrievalAugmentationAdvisor)
                .stream()
                .content();
    }

    /**
     * 聊天记录
     */
    @GetMapping(value = "/messages")
    @SneakyThrows
    @ResponseBody
    public List<Message> messages(@RequestParam(name="conversationId", defaultValue = "feng123") String conversationId) {
        return chatMemory.get(conversationId);
    }

    /**
     * 清空聊天记录
     */
    @GetMapping(value = "/clearMessages")
    @SneakyThrows
    @ResponseBody
    public String clearMessages(@RequestParam(name="conversationId", defaultValue = "feng123") String conversationId) {
        chatMemory.clear(conversationId);
        return "success";
    }

    /**
     * 向量数据库增加数据
     */
    @ResponseBody
    @SneakyThrows
    @PostMapping(value = "/addDocuments")
    public String addDocuments(@RequestParam(name = "documents") List<String> documents) {
        List<Document> collected = documents.stream().map(Document::new).collect(Collectors.toList());
        vectorStore.add(collected);
        return "success";
    }

    @SuppressWarnings("all")
    private static class SelectedFirstDocumentPostProcessor implements DocumentPostProcessor {
        @Override
        public List<Document> process(Query query, List<Document> documents) {
            if (documents.isEmpty()) {
                return Collections.emptyList();
            }
            return Collections.singletonList(documents.get(0));
        }
    }
}
